package com.shuigu.ai.handler;

import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.shuigu.ai.constant.AppConstant;
import com.shuigu.ai.core.builder.VueProjectBuilder;
import com.shuigu.ai.model.enums.StreamMessageTypeEnum;
import com.shuigu.ai.model.message.*;
import com.shuigu.ai.tools.BaseTool;
import com.shuigu.ai.tools.ToolManager;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.memory.ChatMemory;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Flux;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Function;

/**
 * JSON 消息流处理器
 * 处理 VUE_PROJECT 类型的复杂流式响应，包含工具调用信息
 */
@Slf4j
@Component
public class JsonMessageStreamHandler {

    @Resource
    private VueProjectBuilder vueProjectBuilder;

    @Resource
    private ToolManager toolManager;

    /**
     * 处理 TokenStream（VUE_PROJECT）
     * 解析 JSON 消息并重组为完整的响应格式
     *
     * @param originFlux         原始流
     * @param appId              应用ID
     * @return 处理后的流
     */
    public Flux<String> handle(Flux<String> originFlux, ChatMemory chatMemory, long appId) {
        StringBuffer chatHistory = new StringBuffer();
        Set<String> seenToolIds = ConcurrentHashMap.newKeySet();
        // 缓存工具消息（TOOL_REQUEST / TOOL_EXECUTED）要输出的内容
        List<String> pendingToolOutputs = Collections.synchronizedList(new ArrayList<>());

        // helper: 快速解析类型（为了代码清晰，这里做一次轻解析）
        Function<String, StreamMessageTypeEnum> detectType = chunk -> {
            StreamMessage sm = JSONUtil.toBean(chunk, StreamMessage.class);
            return StreamMessageTypeEnum.getEnumByValue(sm.getType());
        };

        // 思考-AI回复文本-工具执行请求-工具执行结果
        return originFlux
                .flatMap(chunk -> {
                    // 对每个 chunk 做处理，map → 可能返回多个输出（工具消息 flush + 本 chunk 输出）
                    StreamMessageTypeEnum t = detectType.apply(chunk);
                    String chunkOut = handleJsonMessageChunk(chunk, chatHistory, seenToolIds);

                    // 创建一个 Flux<String>，这个 chunk 的输出路径可能包括之前缓存 + 当前 chunk 本身
                    List<String> outs = new ArrayList<>();

                    switch (t) {
                        case THINK, AI_RESPONSE -> {
                            // 在输出 AI / 思考之前，把 pendingToolOutputs 先 flush 出来
                            synchronized (pendingToolOutputs) {
                                if (!pendingToolOutputs.isEmpty()) {
                                    outs.addAll(new ArrayList<>(pendingToolOutputs));
                                    pendingToolOutputs.clear();
                                }
                            }
                            if (StrUtil.isNotEmpty(chunkOut)) {
                                outs.add(chunkOut);
                            }
                        }
                        case TOOL_REQUEST, TOOL_EXECUTED -> {
                            // 不立刻输出，而是缓存
                            if (StrUtil.isNotEmpty(chunkOut)) {
                                pendingToolOutputs.add(chunkOut);
                            }
                            // 这一 chunk 本身不产生即时输出（或者可以根据设计决定是否也要输出一部分）
                        }
                        default -> {
                            // 不支持类型，忽略或日志
                            if (StrUtil.isNotEmpty(chunkOut)) {
                                outs.add(chunkOut);
                            }
                        }
                    }

                    return Flux.fromIterable(outs);
                })
                .doOnComplete(() -> {
                    // 最后 flush 一次 pendingToolOutputs（若还有没输出的工具消息）
                    synchronized (pendingToolOutputs) {
                        for (String msg : pendingToolOutputs) {
                            // 这里用 side-effect 输出 —— 也可以改为把它们 emit 出去
                            // 但因为 doOnComplete 里无法发 next，我们可以把这些缓存在 chatMemory 或另触发一次小 Flux
                            // 简单起见，这里把它们写到 chatHistory，并且你可能需要另做手段把它们发给前端
                            chatHistory.append(msg);
                        }
                    }

                    // 原有的完成逻辑
                    String aiResponse = chatHistory.toString();
                    chatMemory.add(AiMessage.aiMessage(aiResponse));
                    String projectPath = AppConstant.CODE_OUTPUT_ROOT_DIR + "/vue_project_" + appId;
                    vueProjectBuilder.buildProjectAsync(projectPath);
                })
                .doOnError(error -> {
                    String errMsg = "AI回复失败: " + error.getMessage();
                    chatMemory.add(AiMessage.aiMessage(errMsg));
                });
    }

    /**
     * 解析并收集 TokenStream 数据
     */
    private String handleJsonMessageChunk(String chunk, StringBuffer chatHistoryStringBuffer, Set<String> seenToolIds) {
        // 解析 JSON
        StreamMessage streamMessage = JSONUtil.toBean(chunk, StreamMessage.class);
        StreamMessageTypeEnum typeEnum = StreamMessageTypeEnum.getEnumByValue(streamMessage.getType());
        switch (typeEnum) {
            case THINK -> {
                ThinkingMessage thinkingMessage = JSONUtil.toBean(chunk, ThinkingMessage.class);
                chatHistoryStringBuffer.append(thinkingMessage.getText());
                return thinkingMessage.getText();
            }
            case AI_RESPONSE -> {
                AiResponseMessage aiMessage = JSONUtil.toBean(chunk, AiResponseMessage.class);
                String data = aiMessage.getData();
                // 直接拼接响应
                chatHistoryStringBuffer.append(data);
                return data;
            }
            case TOOL_REQUEST -> {
                ToolRequestMessage toolRequestMessage = JSONUtil.toBean(chunk, ToolRequestMessage.class);
                String toolId = toolRequestMessage.getId();
                String toolName = toolRequestMessage.getName();
                // 检查是否是第一次看到这个工具 ID
                if (toolId != null && !seenToolIds.contains(toolId)) {
                    // 第一次调用这个工具，记录 ID 并返回工具信息
                    seenToolIds.add(toolId);
                    // 根据工具名称获取工具实例
                    BaseTool tool = toolManager.getTool(toolName);
                    chatHistoryStringBuffer.append(tool.generateToolRequestResponse());
                    // 返回格式化的工具调用信息
                    return tool.generateToolRequestResponse();
                } else {
                    // 不是第一次调用这个工具，直接返回空
                    return "";
                }
            }
            case TOOL_EXECUTED -> {
                ToolExecutedMessage toolExecutedMessage = JSONUtil.toBean(chunk, ToolExecutedMessage.class);
                String toolName = toolExecutedMessage.getName();
                JSONObject jsonObject = JSONUtil.parseObj(toolExecutedMessage.getArguments());
                // 根据工具名称获取工具实例并生成相应的结果格式
                BaseTool tool = toolManager.getTool(toolName);
                String result = tool.generateToolExecutedResult(jsonObject);
                // 输出前端和要持久化的内容
                String output = String.format("\n%s\n", result);
                chatHistoryStringBuffer.append(output);
                return output;
            }
            default -> {
                log.error("不支持的消息类型: {}", typeEnum);
                return "";
            }
        }
    }

    /**
     * 处理 TokenStream（VUE_PROJECT）
     * 解析 JSON 消息并重组为完整的响应格式
     *
     * @param originFlux         原始流
     * @param appId              应用ID
     * @return 处理后的流
     */
//    public Flux<String> handle(Flux<String> originFlux, ChatMemory chatMemory, long appId) {
//        // 收集数据用于生成后端记忆格式
//        StringBuilder chatHistoryStringBuilder = new StringBuilder();
//        // 用于跟踪已经见过的工具ID，判断是否是第一次调用
//        Set<String> seenToolIds = new HashSet<>();
//        return originFlux
//                .map(chunk -> {
//                    // 解析每个 JSON 消息块
//                    return handleJsonMessageChunk(chunk, chatHistoryStringBuilder, seenToolIds);
//                })
//                .filter(StrUtil::isNotEmpty) // 过滤空字串
//                .doOnComplete(() -> {
//                    // 流式响应完成后，添加 AI 消息到对话历史
//                    String aiResponse = chatHistoryStringBuilder.toString();
//                    AiMessage aiMessage = AiMessage.aiMessage(aiResponse);
//                    chatMemory.add(aiMessage);
//
//                    // 异步构造 Vue 项目
//                    String projectPath = AppConstant.CODE_OUTPUT_ROOT_DIR + "/vue_project_" + appId;
//                    vueProjectBuilder.buildProjectAsync(projectPath);
//                })
//                .doOnError(error -> {
//                    // 如果AI回复失败，也要记录错误消息
//                    String errorMessage = "AI回复失败: " + error.getMessage();
//                    chatMemory.add(AiMessage.aiMessage(errorMessage));
//                });
//    }

    /**
     * 解析并收集 TokenStream 数据
     */
//    private String handleJsonMessageChunk(String chunk, StringBuilder chatHistoryStringBuilder, Set<String> seenToolIds) {
//        // 解析 JSON
//        StreamMessage streamMessage = JSONUtil.toBean(chunk, StreamMessage.class);
//        StreamMessageTypeEnum typeEnum = StreamMessageTypeEnum.getEnumByValue(streamMessage.getType());
//        switch (typeEnum) {
//            case THINK -> {
//                ThinkingMessage thinkingMessage = JSONUtil.toBean(chunk, ThinkingMessage.class);
//                chatHistoryStringBuilder.append(thinkingMessage.getText());
//                return thinkingMessage.getText();
//            }
//            case AI_RESPONSE -> {
//                AiResponseMessage aiMessage = JSONUtil.toBean(chunk, AiResponseMessage.class);
//                String data = aiMessage.getData();
//                // 直接拼接响应
//                chatHistoryStringBuilder.append(data);
//                return data;
//            }
//            case TOOL_REQUEST -> {
//                ToolRequestMessage toolRequestMessage = JSONUtil.toBean(chunk, ToolRequestMessage.class);
//                String toolId = toolRequestMessage.getId();
//                String toolName = toolRequestMessage.getName();
//                // 检查是否是第一次看到这个工具 ID
//                if (toolId != null && !seenToolIds.contains(toolId)) {
//                    // 第一次调用这个工具，记录 ID 并返回工具信息
//                    seenToolIds.add(toolId);
//                    // 根据工具名称获取工具实例
//                    BaseTool tool = toolManager.getTool(toolName);
//                    // 返回格式化的工具调用信息
//                    return tool.generateToolRequestResponse();
//                } else {
//                    // 不是第一次调用这个工具，直接返回空
//                    return "";
//                }
//            }
//            case TOOL_EXECUTED -> {
//                ToolExecutedMessage toolExecutedMessage = JSONUtil.toBean(chunk, ToolExecutedMessage.class);
//                String toolName = toolExecutedMessage.getName();
//                JSONObject jsonObject = JSONUtil.parseObj(toolExecutedMessage.getArguments());
//                // 根据工具名称获取工具实例并生成相应的结果格式
//                BaseTool tool = toolManager.getTool(toolName);
//                String result = tool.generateToolExecutedResult(jsonObject);
//                // 输出前端和要持久化的内容
//                String output = String.format("\n%s\n", result);
//                chatHistoryStringBuilder.append(output);
//                return output;
//            }
//            default -> {
//                log.error("不支持的消息类型: {}", typeEnum);
//                return "";
//            }
//        }
//    }
}


